# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import functools

from collections import namedtuple
from collections import deque

from coin.base.datetime_util import to_datetime, to_timestamp_int
from coin.exchange.base.order_gateway import OrderSide
from coin.base.book.types import TradeSide

DT_FORMAT = '%Y-%m-%dT%H:%M:%S.%f'


@functools.lru_cache(maxsize=None)
def load_intervals(filepath):
  intervals = []
  with open(filepath, 'r') as f:
    for line in f.readlines():
      begin_str, end_str = line.strip().split(',')
      begin = datetime.datetime.strptime(begin_str, DT_FORMAT)
      end = datetime.datetime.strptime(end_str, DT_FORMAT)
      intervals.append((to_timestamp_int(begin), to_timestamp_int(end)))

  return intervals


class UnavailabilityChecker(object):
  def __init__(self, unavailable_intervals=None, filepath=None):
    if filepath is not None:
      assert unavailable_intervals is None
      unavailable_intervals = load_intervals(filepath)
    self._intervals = unavailable_intervals
    self._last_idx = None

  def _find_interval(self, timestamp):
    if self._last_idx is not None:
      if ((self._last_idx == 0 or self._intervals[self._last_idx - 1][0] <= timestamp)
          and (self._last_idx == len(self._intervals)
               or timestamp < self._intervals[self._last_idx][0])):
        return self._last_idx

    first = 0
    count = len(self._intervals)
    while count > 0:
      step = count // 2
      if self._intervals[first + step][0] <= timestamp:
        first = first + step + 1
        count -= step + 1
      else:
        count = step

    self._last_idx = first
    return first

  def is_unavailable(self, timestamp):
    idx = self._find_interval(timestamp) - 1
    if idx < 0 or len(self._intervals) <= idx:
      return False

    begin_ts, end_ts = self._intervals[idx]
    return begin_ts <= timestamp < end_ts

  @staticmethod
  @functools.lru_cache(maxsize=None)
  def get_from_file(filepath):
    if filepath is None:
      return None
    return UnavailabilityChecker(filepath=filepath)


class SimpleMatcher(object):
  _Fill = namedtuple('_Fill', ['price', 'qty', 'side', 'maker'])
  _Order = namedtuple('_Order', ['price', 'qty', 'timestamp'])

  def __init__(self,
               side,
               execution_delay=0,
               post_only=False,
               ignore_book_fill=False,
               unavail_checker=None):
    self.side = side
    self.execution_delay = execution_delay
    self._post_only = post_only
    self._ignore_book_fill = ignore_book_fill
    self._unavail_checker = unavail_checker

    self.price = 0
    self.qty = 0
    self._order_queue = deque()
    self._last_ask0 = None
    self._last_bid0 = None

  @property
  def dead(self):
    return self.qty == 0 and not self._order_queue

  @property
  def on_the_fly(self):
    return (self.qty == 0 and self._order_queue and self._order_queue[0].qty != 0)

  @property
  def on_the_fly_qty(self):
    assert self.on_the_fly
    return self._order_queue[0].qty

  def set_order(self, price, qty, timestamp):
    if (self._unavail_checker and self._unavail_checker.is_unavailable(timestamp)):
      return
    order = self._Order(price, qty, timestamp)
    self._order_queue.append(order)

  def cancel(self, timestamp):
    if (self._unavail_checker and self._unavail_checker.is_unavailable(timestamp)):
      return None

    if self._order_queue and self._order_queue[-1].qty == 0:
      return None

    tspq = None
    if len(self._order_queue) > 0:
      tspq = (timestamp, self.side, self._order_queue[-1].price, self._order_queue[-1].qty)

    order = self._Order(0, 0, timestamp)
    self._order_queue.append(order)
    return tspq

  def make_fill(self, fill_qty, maker, fill_price=None):
    assert fill_qty <= self.qty
    if fill_qty == 0:
      return None
    if fill_price is None:
      fill_price = self.price
    else:
      assert not maker
    self.qty -= fill_qty
    return self._Fill(fill_price, fill_qty, self.side, maker)

  def _update_time(self, timestamp):
    if self._last_ask0 is None or self._last_bid0 is None:
      return []

    fills = []
    while (self._order_queue
           and self._order_queue[0].timestamp + self.execution_delay <= timestamp):
      self.price = self._order_queue[0].price
      self.qty = self._order_queue[0].qty
      self._order_queue.popleft()

      if self.qty == 0:
        continue

      if self._post_only:
        last_midp = (self._last_ask0 + self._last_bid0) / 2.0
        if ((self.side == OrderSide.BUY and self.price >= last_midp)
            or (self.side == OrderSide.SELL and self.price <= last_midp)):
          self.qty = 0  # Auto-canceled
        continue

      # Taker fill
      if self.side == OrderSide.BUY and self.price >= self._last_ask0:
        # TODO(jaewon): Make it more realistic
        #               * Use book qty to make partial fill
        fills.append(self.make_fill(self.qty, False, fill_price=self._last_ask0))
      elif self.side == OrderSide.SELL and self.price <= self._last_bid0:
        fills.append(self.make_fill(self.qty, False, fill_price=self._last_bid0))

    return fills

  def update_book(self, book):
    fills = self._update_time(book.timestamp)

    if hasattr(book, "_original_book"):
      self._last_ask0 = book._original_book.ask0().price
      self._last_bid0 = book._original_book.bid0().price
    else:
      self._last_ask0 = book.ask0().price
      self._last_bid0 = book.bid0().price

    if not self._ignore_book_fill and self.qty > 0:
      if self.side == OrderSide.BUY and self.price > self._last_ask0:
        fills.append(self.make_fill(self.qty, True))
      elif self.side == OrderSide.SELL and self.price < self._last_bid0:
        fills.append(self.make_fill(self.qty, True))

    return fills

  def update_trade(self, trade_price, trade_qty, trade_side, timestamp):
    fills = self._update_time(timestamp)

    fill_qty = min(trade_qty, self.qty)
    if fill_qty == 0:
      return fills

    if self.side == OrderSide.BUY and trade_side == TradeSide.TRADE_SELL_SIDE:
      if self.price > self._last_bid0:
        if self.price >= trade_price:
          fills.append(self.make_fill(fill_qty, True))
      else:
        if self.price > trade_price:
          fills.append(self.make_fill(fill_qty, True))

    elif self.side == OrderSide.SELL and trade_side == TradeSide.TRADE_BUY_SIDE:
      if self.price < self._last_ask0:
        if self.price <= trade_price:
          fills.append(self.make_fill(fill_qty, True))
      else:
        if self.price < trade_price:
          fills.append(self.make_fill(fill_qty, True))

    return fills


class PositionCalculator(object):
  def __init__(self, maker_fee, taker_fee):
    self.maker_fee = maker_fee
    self.taker_fee = taker_fee

    self.pos = 0
    self.realized_pnl = 0
    self.unrealized_pnl = 0
    self.num_trade = 0
    self.num_buy = 0
    self.num_sell = 0
    self.qty_sum = 0
    self.pq_sum = 0
    self.fee = 0
    self.submit_qty_sum = 0

  @property
  def return_bp(self):
    if self.pq_sum == 0:
      return None
    return self.pnl / self.pq_sum * 10000.

  def update_book(self, book):
    self._last_ask0 = book.ask0().price
    self._last_bid0 = book.bid0().price
    self._update_unrealized_pnl()

  def _update_unrealized_pnl(self):
    mid_price = (self._last_ask0 + self._last_bid0) / 2.
    self.unrealized_pnl = self.pos * mid_price

  def __str__(self):
    return ('[pnl:%.2f=%.2f+%.2f, pos:%.4f, num_trade:%d, qty_sum:%.4f, '
            'pq_sum:%.4f, return_bp:%.4f]' % (self.pnl,
                                              self.realized_pnl,
                                              self.unrealized_pnl,
                                              self.pos,
                                              self.num_trade,
                                              self.qty_sum,
                                              self.pq_sum,
                                              self.return_bp))

  @property
  def pnl(self):
    return self.realized_pnl + self.unrealized_pnl

  def submit(self, *, timestamp, price, qty, side):
    self.submit_qty_sum += qty

  def update(self, price, qty, side, maker):
    fee = self.maker_fee if maker else self.taker_fee
    if side == OrderSide.BUY:
      self.pos += qty
      self.realized_pnl -= (price * qty) * (1. + fee)
      self.num_buy += 1
      self.fee += (price * qty) * (fee)
    elif side == OrderSide.SELL:
      self.pos -= qty
      self.realized_pnl += (price * qty) * (1. - fee)
      self.num_sell += 1
      self.fee += (price * qty) * (fee)
    else:
      raise ValueError(side)

    self.num_trade += 1
    self.qty_sum += qty
    self.pq_sum += price * qty
    self._update_unrealized_pnl()

  def update_by_fill_object(self, *fills):
    filled = []
    for fill in fills:
      if fill is not None:
        self.update(fill.price, fill.qty, fill.side, fill.maker)
        filled.append(fill)
    return filled

  def get_summary(self, key_postfix=None):
    key_postfix = key_postfix or ''
    if key_postfix != '' and not key_postfix.startswith('_'):
      key_postfix = '_' + key_postfix

    summary = {
        'pos': self.pos,
        'pnl': self.pnl,
        "fee": self.fee,
        "fee_bp": self.fee / (1e-9 + self.pq_sum) * 1e4,
        'realized_pnl': self.realized_pnl,
        'unrealized_pnl': self.unrealized_pnl,
        'num_trade': self.num_trade,
        'num_buy': self.num_buy,
        'num_sell': self.num_sell,
        'qty_sum': self.qty_sum,
        'pq_sum': self.pq_sum,
        'return_bp': self.return_bp,
        'sub_qty_sum': self.submit_qty_sum,
    }
    return {(key + key_postfix): value for key, value in summary.items()}
